import os
import argparse
import json
from typing import Any

# 从环境变量中获取配置项的值
class Parser(object):
    def __init__(self,
                 name: str, default: Any = None,
                 type_: callable = None, help_info: str = None):
        self.name = name
        self.default = default
        self.type_ = type_
        self.help = help_info

    def __set__(self, instance, value):
        setattr(self, 'value', value)

    def __get__(self, instance, owner):
        if hasattr(self, 'value'):
            return getattr(self, 'value')
        v = os.environ.get(self.name, self.default)
        if self.type_ is not None and self.name in os.environ:
            v = self.type_(v)
        setattr(self, 'value', v)
        return v

    def __call__(self, s):
        if self.type_ is not None:
            return self.type_(s)
        return s

    def __str__(self):
        return f'<{self.name}: {self.help} default:{self.default}>'

    def __repr__(self) -> str:
        return self.__str__()

class Config(object):
    def __init__(self):
        pass

    @property
    def value_dict(self):
        return self._search_cfg_recursively(self.__class__)

    # 递归搜索当前类及其所有父类中的配置项，并将结果以字典的形式返回
    @staticmethod
    def _search_cfg_recursively(root):
        vals = dict()
        for base in root.__bases__:
            vals.update(Config._search_cfg_recursively(base))
        for k, v in root.__dict__.items():
            if isinstance(v, Parser):
                vals[k] = v.__get__(None, None)
        return vals

    # 类的字符串表示方法，打印类的实例时，会调用该方法。
    # 返回一个格式化的字符串，包含类名和配置项的字典表示，使用 json.dumps 方法将字典转换为 JSON 字符串
    def __repr__(self):
        return f'<{self.__class__.__name__}: {json.dumps(self.value_dict)}>'

    def sample_cfg(self):
        for k, v in self.__class__.__dict__.items():
            if isinstance(v, Parser):
                print(f'{k}={v.default} ', end='')
        print()

class TrainerConfig(Config):
    debug = Parser('debug', False,
                   lambda x: not x.lower().startswith('f'), 'debug mode ')
    # /home/zhangyichi/projects/DeepDR_Plus/logs/bs16_epoch150_lr0.0001/TrainerDR_20250409_seemshigh/model_054.pth
    # /home/zhangyichi/projects/DeepDR_Plus/logs/bs16_epoch150_lr3e-05/TrainerDR_20250408_raw/model_149.pth
    # /home/zhangyichi/projects/DeepDR_Plus/logs/bs16_epoch150_lr3e-05/TrainerDR_20250408_big/model_149.pth
    load_pretrain = Parser('load_pretrain', None, 
                           str, 'load pretrained model')
    load_checkpoint = Parser('load_checkpoint', None, 
                             str, 'load checkpoint to continue training')
    bright_adjust = Parser('bright_adjust', False, bool, 'if use images after brightness adjusting')
    batch_size = Parser('batch_size', 16, int, 'batch size')
    start_epoch = Parser('start_epoch', 0, int, 'this will be set when checkpoint is loaded')
    epochs = Parser('epochs', 50, int, 'number of max epochs to train')
    image_size = Parser('image_size', 512, int, 'image size')
    lr = Parser('lr', 1e-4, float, 'learning rate')
    device = Parser('device', 'cuda:0', str, 'device')
    num_workers = Parser('num_workers', 4, int, 'number of workers')
    model = Parser('model', 'resnet101', str, 'backbone model')
    model_path = Parser('model_path', '/home/zhangyichi/projects/DeepDR_Plus/logs/bs16_epoch150_lr3e-05/TrainerDR_20250408_raw/model_149.pth', str, 'load for test')
    result_path = Parser('result_path', 'result_final_0408_raw', str, 'save predict result')